package com.uyayo.framework.limiter.interceptor;

import com.google.common.base.Strings;
import com.google.common.cache.Cache;
import com.google.common.cache.CacheBuilder;
import com.google.common.cache.CacheLoader;
import com.google.common.cache.LoadingCache;
import com.uyayo.framework.limiter.exception.CsrfException;
import com.uyayo.framework.limiter.exception.CsrfExpireException;
import com.uyayo.framework.limiter.exception.RefererException;
import com.uyayo.framework.limiter.support.RequestLimiterHelper;
import com.uyayo.framework.limiter.annotation.RequestLimiter;
import com.uyayo.framework.limiter.constants.CharacterConstants;
import com.uyayo.framework.limiter.exception.RequestException;
import com.uyayo.framework.limiter.utils.RequestUtils;
import io.jsonwebtoken.ExpiredJwtException;
import io.jsonwebtoken.Jwts;
import io.jsonwebtoken.MissingClaimException;
import org.apache.commons.lang3.StringEscapeUtils;
import org.aspectj.lang.JoinPoint;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.web.context.request.RequestContextHolder;
import org.springframework.web.context.request.ServletRequestAttributes;

import javax.servlet.http.HttpServletRequest;
import java.lang.invoke.MethodHandles;
import java.text.MessageFormat;
import java.util.Arrays;
import java.util.List;
import java.util.Optional;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

/**
 * @author hfj
 * @date 2019/5/23
 * csrfTokenCount计数非线程安全拦截器
 */

public class RequestLimiterInterceptor {
    private static final Logger log = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass());
    private final RequestLimiterHelper requestLimiterHelper;
    private static final String LIMITER_KEY = "REQUEST_LIMITER_INTERCEPTOR";
    private static final long JWT_TOKEN_CACHE_DEFAULT_EXPIRATION = 180 * 1000;
    private static final int REQUEST_LIMITER_COUNT = 5;
    private static final String WILDCARD = "*";
    private final Cache<String, Long> JWT_TOKEN_CACHE;
    private final LoadingCache<String, List<String>> WHITELIST;
    private String jwtKey;
    private Integer limiterCount;

    public RequestLimiterInterceptor(RequestLimiterHelper requestLimiterHelper,String requestLimiterWhitelist) {
        this(requestLimiterHelper, RequestLimiterInterceptor.class.getName(), JWT_TOKEN_CACHE_DEFAULT_EXPIRATION,requestLimiterWhitelist);
    }

    public RequestLimiterInterceptor(RequestLimiterHelper requestLimiterHelper, String jwtKey,String requestLimiterWhitelist) {
        this(requestLimiterHelper, jwtKey, JWT_TOKEN_CACHE_DEFAULT_EXPIRATION,requestLimiterWhitelist);
    }

    public RequestLimiterInterceptor(RequestLimiterHelper requestLimiterHelper, String jwtKey, Long expiration,String requestLimiterWhitelist) {
        this.requestLimiterHelper = requestLimiterHelper;
        this.jwtKey = jwtKey;
        JWT_TOKEN_CACHE = CacheBuilder.newBuilder().expireAfterWrite(expiration, TimeUnit.MILLISECONDS).build();
        WHITELIST = CacheBuilder.newBuilder().refreshAfterWrite(5, TimeUnit.MINUTES).build(new CacheLoader<String, List<String>>() {
            @Override
            public List<String> load(String key) throws Exception {
                if(Strings.isNullOrEmpty(requestLimiterWhitelist)){
                    return Arrays.stream(new String[]{""}).collect(Collectors.toList());
                }
                    String[] requestLimiterWhitelistArr = requestLimiterWhitelist.split(",");
                    return Arrays.stream(requestLimiterWhitelistArr).collect(Collectors.toList());
                }
        });
    }

    public void limiterMethod(JoinPoint joinPoint, RequestLimiter limiter) {
        //RequestInfo requestInfo = null;
        HttpServletRequest request = null;
        Object[] objects = joinPoint.getArgs();
        for (Object o : objects) {
            //if (o instanceof RequestInfo) {
            //    requestInfo = (RequestInfo) o;
            //}
            if (o instanceof HttpServletRequest) {
                request = (HttpServletRequest) o;
            }
        }
        if (request == null) {
            request = ((ServletRequestAttributes) (RequestContextHolder.currentRequestAttributes())).getRequest();
            if(request==null){
                throw new IllegalArgumentException("方法内必须要包含HttpServletRequest");
            }
        }
        //if (requestInfo == null) {
        //    requestInfo = RequestUtils.getRequestInfo(request);
        //}
        final String ip = RequestUtils.getIpAddr(request);
        try {
            if (limiter.referer()) {
                String referer = request.getHeader("Referer");
                if (Strings.isNullOrEmpty(referer)) {
                    throw new RefererException(RequestUtils.getIpAddr(request) + "请求时referer为空!");
                }
            }
            String csrfTokenKey = limiter.csrfToken();
            if (!Strings.isNullOrEmpty(csrfTokenKey)) {
                String token = request.getHeader(csrfTokenKey);
                if (Strings.isNullOrEmpty(token)) {
                    token = request.getParameter(csrfTokenKey);
                }
                if (Strings.isNullOrEmpty(token)) {
                    throw new CsrfException(RequestUtils.getIpAddr(request) + "请求时csrfToken为空!");
                }
                token = StringEscapeUtils.unescapeHtml4(token);
                Long count = JWT_TOKEN_CACHE.getIfPresent(token);
                if (count != null && limiter.tokenCount() >= count) {
                    throw new CsrfException(MessageFormat.format("{0}请求时token已超过使用次数!count:{1},limiter:{2}", ip, count, limiter.tokenCount()));
                }
                try {
                    Jwts.parser().requireSubject("ag")
                            .require(CharacterConstants.REQUEST_IP_STRING, RequestUtils.getIpAddr(request))
                            .setSigningKey(getJwtKey()).parse(token);
                } catch (MissingClaimException e) {
                    throw new CsrfException(RequestUtils.getIpAddr(request) + "请求时IP不符合!" + e.getMessage());
                } catch (ExpiredJwtException e) {
                    throw new CsrfExpireException(RequestUtils.getIpAddr(request) + "请求时token已过期!" + e.getMessage());
                } catch (Exception e) {
                    throw new CsrfException(RequestUtils.getIpAddr(request) + "请求时jwtToken检查不通过!" + e.getMessage());
                }
                JWT_TOKEN_CACHE.put(token, count == null ? 1 : count + 1);
            }
            List<String> whiteList = WHITELIST.getUnchecked("whitelist");
            Optional optional = whiteList.stream().filter(s -> ipMatch(s, ip)).findAny();
            if (!optional.isPresent()) {
                int count;
                if (limiter.count() != 0) {
                    count = limiter.count();
                } else if (null != limiterCount) {
                    count = limiterCount;
                } else {
                    count = REQUEST_LIMITER_COUNT;
                }
                if (count > 0) {
                    boolean ipParallel = limiter.ipParallel();
                    String[] strings = limiter.params();
                    StringBuilder limiterKey = new StringBuilder(LIMITER_KEY).append(CharacterConstants.COLON);
                    if (!Strings.isNullOrEmpty(limiter.prefix())) {
                        limiterKey.append(limiter.prefix()).append(CharacterConstants.COLON);
                    }
                    if(ipParallel){
                        //联立IP进行限制
                        limiterKey.append(ip);
                        for (String param : strings) {
                            limiterKey.append(CharacterConstants.COLON).append(String.valueOf(request.getParameter(param)));
                        }
                        requestLimiterHelper.incrRequestLimiter(limiterKey.toString(), limiter.seconds(), count);

                    }else{
                        //不联立IP进行限制
                        String ipLimiterKey = LIMITER_KEY + CharacterConstants.COLON +
                                ip;
                        //先限制IP
                        requestLimiterHelper.incrRequestLimiter(ipLimiterKey, limiter.seconds(), count);


                        //再限制其他参数
                        for (String param : strings) {
                            limiterKey.append(CharacterConstants.COLON).append(String.valueOf(request.getParameter(param)));
                        }
                        requestLimiterHelper.incrRequestLimiter(limiterKey.toString(), limiter.seconds(), count);
                    }

                }
            }
        } catch (Exception e) {
            if (limiter.exception().length > 0) {
                throw getException(limiter.exception()[0], e.getMessage());
            } else {
                throw e;
            }
        }
    }

    private boolean ipMatch(String white, String ip) {
        if (!white.endsWith(WILDCARD)) {
            return white.equalsIgnoreCase(ip);
        }
        return ip.startsWith(white.substring(0, white.length() - 1));
    }

    public Integer getLimiterCount() {
        return limiterCount;
    }

    public void setLimiterCount(Integer limiterCount) {
        this.limiterCount = limiterCount;
    }

    private String getJwtKey() {
        return jwtKey;
    }

    public void setJwtKey(String jwtKey) {
        this.jwtKey = jwtKey;
    }

    private <T extends RequestException> T getException(Class<T> t, Object o) {
        try {
            T exception = t.newInstance();
            exception.setObject(o);
            return exception;
        } catch (Exception e) {
            log.error("" + e);
            throw new RuntimeException(e);
        }
    }


}
